#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
监控爬虫进度的脚本
"""

import os
import sqlite3
import time
from datetime import datetime

def check_progress(db_path="./output/index.db"):
    """检查爬取进度
    
    Args:
        db_path: 数据库路径
    """
    try:
        # 连接数据库
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 获取资源总数
        cursor.execute("SELECT COUNT(*) FROM resources")
        total_resources = cursor.fetchone()[0]
        
        # 获取各类型资源数量
        cursor.execute("SELECT type, COUNT(*) FROM resources GROUP BY type")
        type_counts = dict(cursor.fetchall())
        
        # 获取最近爬取的资源
        cursor.execute(
            "SELECT url, type, crawl_time FROM resources ORDER BY crawl_time DESC LIMIT 5"
        )
        recent_resources = cursor.fetchall()
        
        # 获取队列中待处理的URL数量
        cursor.execute("SELECT COUNT(*) FROM url_queue WHERE status = 'pending'")
        pending_urls = cursor.fetchone()[0]
        
        # 获取已处理的URL数量
        cursor.execute("SELECT COUNT(*) FROM url_queue WHERE status = 'done'")
        done_urls = cursor.fetchone()[0]
        
        # 获取处理失败的URL数量
        cursor.execute("SELECT COUNT(*) FROM url_queue WHERE status = 'error'")
        error_urls = cursor.fetchone()[0]
        
        # 关闭连接
        conn.close()
        
        # 输出进度报告
        print("\n===== 爬虫进度报告 =====")
        print(f"报告时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"已爬取资源总数: {total_resources}")
        print("\n资源类型分布:")
        for type_name, count in type_counts.items():
            print(f"  - {type_name}: {count}")
        
        print("\nURL队列状态:")
        print(f"  - 待处理: {pending_urls}")
        print(f"  - 已完成: {done_urls}")
        print(f"  - 失败: {error_urls}")
        
        print("\n最近爬取的资源:")
        for url, type_name, crawl_time in recent_resources:
            time_str = datetime.fromisoformat(crawl_time).strftime('%Y-%m-%d %H:%M:%S')
            print(f"  - [{type_name}] {url} ({time_str})")
        
        print("\n输出目录大小:")
        output_size = get_dir_size("./output")
        print(f"  - 总大小: {format_size(output_size)}")
        
        print("========================\n")
        
    except sqlite3.Error as e:
        print(f"数据库错误: {e}")
    except Exception as e:
        print(f"检查进度时出错: {e}")

def get_dir_size(path):
    """获取目录大小
    
    Args:
        path: 目录路径
        
    Returns:
        int: 目录大小（字节）
    """
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(path):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            if not os.path.islink(fp):
                total_size += os.path.getsize(fp)
    return total_size

def format_size(size_bytes):
    """格式化文件大小
    
    Args:
        size_bytes: 文件大小（字节）
        
    Returns:
        str: 格式化后的文件大小
    """
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_bytes < 1024.0 or unit == 'TB':
            break
        size_bytes /= 1024.0
    return f"{size_bytes:.2f} {unit}"

if __name__ == "__main__":
    check_progress()
